import streamlit as st
import os
from rag_system import LocalRAGSystem

# 页面配置
st.set_page_config(
    page_title="本地知识库问答系统",
    page_icon="🤖",
    layout="wide"
)

# 初始化会话状态
if 'rag_system' not in st.session_state:
    st.session_state.rag_system = None
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

# 侧边栏配置
with st.sidebar:
    st.title("系统配置")
    
    # API密钥输入
    api_key = st.text_input("硅基流动API密钥", type="password")
    if api_key:
        os.environ["GUIJI_API_KEY"] = api_key
    
    # 文档上传
    st.subheader("文档管理")
    uploaded_files = st.file_uploader(
        "上传文档", 
        accept_multiple_files=True,
        type=['txt']
    )
    
    if uploaded_files and st.button("加载文档"):
        # 创建临时文档目录
        if not os.path.exists("temp_docs"):
            os.makedirs("temp_docs")
            
        # 保存上传的文件
        for file in uploaded_files:
            with open(f"temp_docs/{file.name}", "wb") as f:
                f.write(file.getvalue())
        
        # 初始化RAG系统
        st.session_state.rag_system = LocalRAGSystem()
        st.session_state.rag_system.load_documents("temp_docs")
        st.success("文档加载成功！")

# 主界面
st.title("🤖 本地知识库问答系统")

# 显示聊天历史
for query, response in st.session_state.chat_history:
    with st.chat_message("user"):
        st.write(query)
    with st.chat_message("assistant"):
        st.write(response["answer"])
        with st.expander("查看参考文档"):
            for doc in response["source_documents"]:
                st.markdown(f"""
                **文档片段：**
                ```
                {doc.page_content[:200]}...
                ```
                **来源：** {doc.metadata}
                ---
                """)

# 用户输入
user_input = st.chat_input("请输入您的问题")

if user_input:
    if not st.session_state.rag_system:
        st.error("请先上传并加载文档！")
    elif not os.getenv("GUIJI_API_KEY"):
        st.error("请先配置硅基流动API密钥！")
    else:
        # 显示用户问题
        with st.chat_message("user"):
            st.write(user_input)
            
        # 显示系统响应
        with st.chat_message("assistant"):
            with st.spinner("思考中..."):
                result = st.session_state.rag_system.query(
                    user_input, 
                    [(q, r["answer"]) for q, r in st.session_state.chat_history]
                )
                
                st.write(result["answer"])
                with st.expander("查看参考文档"):
                    for doc in result["source_documents"]:
                        st.markdown(f"""
                        **文档片段：**
                        ```
                        {doc.page_content[:200]}...
                        ```
                        **来源：** {doc.metadata}
                        ---
                        """)
        
        # 更新聊天历史
        st.session_state.chat_history.append((user_input, result)) 